package com.liyunc.demo;

import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacpp.BytePointer;
import org.bytedeco.javacpp.Pointer;
import org.bytedeco.javacpp.PointerPointer;
import org.bytedeco.javacpp.avcodec;
import org.bytedeco.javacpp.avformat;
import org.bytedeco.javacpp.avutil;
import org.bytedeco.javacpp.swscale;

import java.io.File;
import java.util.Arrays;

import static org.bytedeco.javacpp.avcodec.av_new_packet;
import static org.bytedeco.javacpp.avcodec.av_packet_unref;
import static org.bytedeco.javacpp.avcodec.avcodec_alloc_context3;
import static org.bytedeco.javacpp.avcodec.avcodec_close;
import static org.bytedeco.javacpp.avcodec.avcodec_find_encoder;
import static org.bytedeco.javacpp.avcodec.avcodec_open2;
import static org.bytedeco.javacpp.avcodec.avcodec_parameters_to_context;
import static org.bytedeco.javacpp.avcodec.avcodec_receive_packet;
import static org.bytedeco.javacpp.avcodec.avcodec_send_frame;
import static org.bytedeco.javacpp.avformat.AVIO_FLAG_READ_WRITE;
import static org.bytedeco.javacpp.avformat.av_dump_format;
import static org.bytedeco.javacpp.avformat.av_guess_format;
import static org.bytedeco.javacpp.avformat.av_read_frame;
import static org.bytedeco.javacpp.avformat.av_write_frame;
import static org.bytedeco.javacpp.avformat.av_write_trailer;
import static org.bytedeco.javacpp.avformat.avformat_alloc_context;
import static org.bytedeco.javacpp.avformat.avformat_close_input;
import static org.bytedeco.javacpp.avformat.avformat_free_context;
import static org.bytedeco.javacpp.avformat.avformat_new_stream;
import static org.bytedeco.javacpp.avformat.avformat_write_header;
import static org.bytedeco.javacpp.avformat.avio_close;
import static org.bytedeco.javacpp.avformat.avio_open;
import static org.bytedeco.javacpp.avutil.AVMEDIA_TYPE_VIDEO;
import static org.bytedeco.javacpp.avutil.AV_LOG_ERROR;
import static org.bytedeco.javacpp.avutil.AV_PIX_FMT_YUVJ420P;
import static org.bytedeco.javacpp.avutil.av_frame_alloc;
import static org.bytedeco.javacpp.avutil.av_image_fill_arrays;
import static org.bytedeco.javacpp.avutil.av_image_get_buffer_size;
import static org.bytedeco.javacpp.avutil.av_log_set_level;
import static org.bytedeco.javacpp.avutil.av_malloc;
import static org.bytedeco.javacpp.swscale.SWS_BICUBIC;
import static org.bytedeco.javacpp.swscale.sws_freeContext;
import static org.bytedeco.javacpp.swscale.sws_getContext;
import static org.bytedeco.javacpp.swscale.sws_scale;

/**
 * 参考连接: https://blog.csdn.net/boling_cavalry/article/details/119062543
 * 将视频转化为图片.
 */
@Slf4j
public class MediaToImage {

    private static final int DEN = 25;
    private static final int LINE = 3;

    /**
     * 打开流媒体，取一帧，转为YUVJ420P，在保存为jpg文件.
     */
    public void openMediaAndSaveImage(String url, String targetFile) {
        log.info("is opening file [{}]", url);

        // 打开指定流媒体，进行解封装，得到解封上下文
        avformat.AVFormatContext formatContext = getFormatContext(url);
        if (null == formatContext) {
            log.error("get formatContext fail");
            return;
        }

        // 控制台打印流媒体信息
        avformat.av_dump_format(formatContext, 0, url, 0);

        // 流媒体解封装后有一个保存了所有流的数组，videoStreamIndex表示视频流在数组中的位置
        int videoStreamIndex = getVideoStreamIndex(formatContext);
        // 找不到视频流就直接返回
        if (videoStreamIndex < 0) {
            log.error("not found video stream");
            return;
        }

        log.info("video stream start from [{}]", videoStreamIndex);

        // 得到解码上下文，已经完成了初始化
        avcodec.AVCodecContext codecContext = getCodecContext(formatContext, videoStreamIndex);
        if (null == codecContext) {
            log.error("create codecContext fail");
            return;
        }

        // 从视频流中解码第一帧
        avutil.AVFrame firstFrame = getSingleFrame(codecContext, formatContext, videoStreamIndex);
        if (null == firstFrame) {
            log.error("get frame from video stream fail");
            return;
        }

        // 将YUV420P图像转换YUVJ420P
        // 转换后的图片的AVFrame，及其对应的数据指针，都放在frameData对象中
        FrameData frameData = YUV420PToYUVJ420P(codecContext, firstFrame);
        if (null == frameData) {
            log.error("YUV420P to YUVJ420P fail");
            return;
        }

        // 持久化存储
        saveImage(frameData.avFrame, targetFile);

        // 按顺序释放
        release(true,
            null,
            null,
            codecContext,
            formatContext,
            frameData.buffer,
            frameData.avFrame,
            firstFrame);

        log.info("success");
    }

    /**
     * 释放资源，顺序是先释放数据，再释放上下文.
     */
    private void release(boolean isInput, avcodec.AVPacket avPacket, avformat.AVIOContext ioContext,
                         avcodec.AVCodecContext avCodecContext,
                         avformat.AVFormatContext avFormatContext, Pointer... pointers) {
        if (null != avPacket) {
            av_packet_unref(avPacket);
        }

        // 解码后，这是个数组，要遍历处理
        if (null != pointers) {
            Arrays.stream(pointers).forEach(avutil::av_free);
        }

        // 解码
        if (null != avCodecContext) {
            avcodec_close(avCodecContext);
        }

        // 解协议
        if (null != ioContext) {
            avio_close(ioContext);
        }

        // 解封装
        if (null != avFormatContext) {
            if (isInput) {
                avformat_close_input(avFormatContext);
            } else {
                avformat_free_context(avFormatContext);
            }
        }
    }

    /**
     * 将传入的帧以图片的形式保存在指定位置.
     *
     * @param avFrame    帧数据
     * @param targetFile 存储位置
     */
    private void saveImage(avutil.AVFrame avFrame, String targetFile) {
        // 设置FFmpeg日志级别（默认是debug，设置成error）
        av_log_set_level(AV_LOG_ERROR);

        avcodec.AVPacket avPacket = null;
        avformat.AVStream avStream;

        int width = avFrame.width();
        int height = avFrame.height();

        // 分配AVFormatContext对象
        avformat.AVFormatContext avFormatContext = avformat_alloc_context();

        // 分配avCodecContext对象
        avcodec.AVCodecContext avCodecContext = avcodec_alloc_context3(null);

        // 设置输出格式（设计到封装和容器）
        avFormatContext.oformat(av_guess_format("mjpeg", null, null));
        if (null == avFormatContext.oformat()) {
            log.error("output format set fail");
            return;
        }

        try {
            // 创建并初始化一个和该url相关的AVIOContext
            avformat.AVIOContext ioContext = new avformat.AVIOContext();

            // 判断target_file路径是否存在
            File file = new File(targetFile);
            File fileParent = file.getParentFile();
            if (!fileParent.exists()) {
                boolean result = fileParent.mkdirs();
                if (!result) {
                    log.error("[{}]创建失败", fileParent);
                    return;
                }
            }

            // 打开输出文件
            if (avio_open(ioContext, targetFile, AVIO_FLAG_READ_WRITE) < 0) {
                // 输出路径不存在会失败
                log.info("output file open fail");
                return;
            }

            // 封装之上是协议，这里将封装上下文和协议上下文关联
            avFormatContext.pb(ioContext);

            // 构建一个新stream
            avStream = avformat_new_stream(avFormatContext, null);
            if (null == avStream) {
                log.error("stream put file fail");
                return;
            }

            int codecId = avFormatContext.oformat().video_codec();

            // 设置该stream的信息
            avcodec_parameters_to_context(avCodecContext, avStream.codecpar());
            avCodecContext.codec_id(codecId);
            avCodecContext.codec_type(AVMEDIA_TYPE_VIDEO);
            avCodecContext.pix_fmt(AV_PIX_FMT_YUVJ420P);
            avCodecContext.width(width);
            avCodecContext.height(height);
            avCodecContext.time_base().num(1);
            avCodecContext.time_base().den(DEN);

            // 打印媒体信息
            av_dump_format(avFormatContext, 0, targetFile, 1);

            // 查找解码器
            avcodec.AVCodec avCodec = avcodec_find_encoder(codecId);
            if (null == avCodec) {
                log.info("get avCodec fail");
                return;
            }

            // 用解码器来初始化解码上下文
            if (avcodec_open2(avCodecContext, avCodec, (PointerPointer<Pointer>) null) < 0) {
                log.error("init avCodecContext fail");
                return;
            }

            // 输出的packet
            avPacket = new avcodec.AVPacket();

            // 分配
            if (av_new_packet(avPacket, width * height * LINE) < 0) {
                return;
            }

            // 把流的头信息写到要输出的媒体文件中
            avformat_write_header(avFormatContext, (PointerPointer<Pointer>) null);

            // 把帧的内容进行编码
            int ret = avcodec_send_frame(avCodecContext, avFrame);
            int get = avcodec_receive_packet(avCodecContext, avPacket);
            if (ret != 0 && get != 0) {
                log.error("frame to packet fail");
                return;
            }

            // 输出第一帧
            if (av_write_frame(avFormatContext, avPacket) < 0) {
                log.error("output fistFrame fail");
                return;
            }

            // 写文件尾
            if (av_write_trailer(avFormatContext) < 0) {
                log.error("write trailer fail");
            }

        } catch (Exception exception) {
            log.error(exception.toString());
        } finally {
            // 资源清理
            release(false, avPacket, avFormatContext.pb(), avCodecContext, avFormatContext);
        }

    }

    /**
     * 将YUV420P格式的图像转为YUVJ420P格式.
     *
     * @param avCodecContext 解码上下文
     * @param avFrame        帧数据
     * @return 转化后的帧极其相对应的数据指针
     */
    private FrameData YUV420PToYUVJ420P(avcodec.AVCodecContext avCodecContext,
                                        avutil.AVFrame avFrame) {
        // 分配一个帧对象，保存从YUV420P转为YUVJ420P的结果
        avutil.AVFrame frameRGB = av_frame_alloc();
        if (null == frameRGB) {
            return null;
        }

        int width = avCodecContext.width();
        int height = avCodecContext.height();

        // 一些参数设定
        frameRGB.width(width);
        frameRGB.height(height);
        frameRGB.format(AV_PIX_FMT_YUVJ420P);

        // 计算转为YUVJ420P之后的图片字节数
        int countBytes = av_image_get_buffer_size(AV_PIX_FMT_YUVJ420P, width, height, 1);

        // 分配内存
        BytePointer buffer = new BytePointer(av_malloc(countBytes));

        // 图片处理工具的初始化操作
        swscale.SwsContext swsContext = sws_getContext(
            width, height,
            avCodecContext.pix_fmt(),
            width, height,
            AV_PIX_FMT_YUVJ420P,
            SWS_BICUBIC,
            null, null, (double[]) null);

        // 将frameRGB的data指针指向刚才分配好的内存(即buffer)
        av_image_fill_arrays(frameRGB.data(), frameRGB.linesize(), buffer, AV_PIX_FMT_YUVJ420P,
            width, height, 1);

        // 转换图像格式，将解压出来的YUV420P的图像转换为YUVJ420P的图像
        sws_scale(swsContext, avFrame.data(), avFrame.linesize(), 0, height, frameRGB.data(),
            frameRGB.linesize());

        // 及时释放
        sws_freeContext(swsContext);

        // 将AVFrame和BytePointer打包到FrameData中返回，这两个对象都要做显示的释放操作
        return new FrameData(frameRGB, buffer);
    }

    /**
     * 获取第一帧后解码.
     *
     * @param avCodecContext   解码上下文
     * @param avFormatContext  解封装上下文
     * @param videoStreamIndex 视频流的位置
     * @return 第一帧数据
     */
    private avutil.AVFrame getSingleFrame(avcodec.AVCodecContext avCodecContext,
                                          avformat.AVFormatContext avFormatContext,
                                          int videoStreamIndex) {
        // 分配帧对象
        avutil.AVFrame avFrame = av_frame_alloc();

        // 是否找到
        boolean exists = false;

        avcodec.AVPacket avPacket = new avcodec.AVPacket();

        try {
            // 每一次循环都会读取一个packet
            while (av_read_frame(avFormatContext, avPacket) >= 0) {
                // 检查packet所属的流是不是视频流
                if (avPacket.stream_index() == videoStreamIndex) {
                    // 将AVPacket解码成AVFrame
                    // 1、第一步 发送编码数据包
                    int ret = avcodec.avcodec_send_packet(avCodecContext, avPacket);
                    // 2、第二部 接收编码数据
                    int get = avcodec.avcodec_receive_frame(avCodecContext, avFrame);

                    if (ret == 0 && get == 0) {
                        exists = true;
                        break;
                    }
                }
            }
            return exists ? avFrame : null;
        } finally {
            // 一定要执行释放操作
            av_packet_unref(avPacket);
        }
    }

    /**
     * 生成解码上下文.
     *
     * @param avFormatContext  解封装上下文
     * @param videoStreamIndex 视频流的位置
     * @return 解码上下文
     */
    private avcodec.AVCodecContext getCodecContext(avformat.AVFormatContext avFormatContext,
                                                   int videoStreamIndex) {
        // 解码器
        avcodec.AVCodec avCodec;

        // 得到解码上下文

        avcodec.AVCodecContext avCodecContext = avcodec_alloc_context3(null);
        avcodec_parameters_to_context(avCodecContext,
            avFormatContext.streams(videoStreamIndex).codecpar());


        // 根据解码上下文得到解码器
        avCodec = avcodec.avcodec_find_decoder(avCodecContext.codec_id());
        if (null == avCodec) {
            return null;
        }

        // 用解码器来初始化解码上下文
        if (avcodec.avcodec_open2(avCodecContext, avCodec, (PointerPointer<Pointer>) null) < 0) {
            return null;
        }

        return avCodecContext;
    }

    /**
     * 流媒体解封装后得到多个流组成的数组，找到视频流在数组中的位置.
     *
     * @param avFormatContext 解封装上下文
     * @return 视频流在数组中的位置
     */
    private int getVideoStreamIndex(avformat.AVFormatContext avFormatContext) {
        int videoStream = -1;

        // 解封转之后有多个流，找出视频流是第几个
        for (int i = 0; i < avFormatContext.nb_streams(); i++) {
            avcodec.AVCodecContext context = avcodec_alloc_context3(null);
            avcodec_parameters_to_context(context, avFormatContext.streams(i).codecpar());
            if (context.codec_type() == AVMEDIA_TYPE_VIDEO) {
                videoStream = i;
                break;
            }
        }
        return videoStream;
    }

    /**
     * 生成解封上下文.
     *
     * @param url 流媒体文件
     * @return 解封装上下文
     */
    private avformat.AVFormatContext getFormatContext(String url) {
        // 解封装上下文
        avformat.AVFormatContext avFormatContext = new avformat.AVFormatContext(null);

        // 打开流媒体文件
        if (avformat.avformat_open_input(avFormatContext, url, null, null) != 0) {
            log.error("open file fail");
            return null;
        }

        // 读取媒流体数据，获得流的信息
        if (avformat.avformat_find_stream_info(avFormatContext, (PointerPointer<Pointer>) null)
            < 0) {
            log.error("get media stream fail");
            return null;
        }
        return avFormatContext;
    }
}
